package com.test.langchain4j.controller;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

import com.test.langchain4j.constant.QdrantConstant;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

/**
 * Created with IntelliJ IDEA.
 *
 * @description:
 * @author: liuziyang
 * @since: 2025/8/4 08:55
 * @modifiedBy:
 * @version: 1.0
 */
@RestController
@Slf4j
public class EmbeddingController {
  @Resource private EmbeddingModel embeddingModel;
  @Resource private QdrantClient qdrantClient;
  @Resource private EmbeddingStore<TextSegment> embeddingStore;

  /**
   * 文本向量化测试，看看形成向量后的文本 http://localhost:9012/embedding/embed
   *
   * @return string
   */
  @GetMapping(value = "/embedding/embed")
  public String embed() {
    String prompt =
        """
        1.铁路行车组织
        organization of train operation
        铁路行车组织是铁路运输组织的重要组成部分，是铁路综合运用各种技术设备合理组织列车运行以实现旅客和货物运输过程的生产计划与组织工作。它的主内容包括：车站技术作业组织、车流组织、列车运行图和铁路通过能力、铁路运输工作技术计划和调度工作等。
        2.铁路技术管理规程
        Regulations Governing Railway Technical Operation
        简称《技规》，是铁路进行技术管理和从事运输生产的基本法规。它的主要内容包括：铁路技术设备的基本要求和标准，行车组织工作应遵循的基本原则、工作方法和作业程序、信号显示的要求和方法以及铁路运输工作人员的主要职责和必须具备的基本条件。
        3.铁路行车组织规则
        Rules Governing Organization of Train Operation；operating rules
        简称《行规》，是铁路局根据《技规》中行车组织篇规定的原则，结合铁路局管内技术设备的具体条件所制定的一种行车组织的补充规则，作为《技规》的补充。
        4.车站行车工作细则
        Detailed Instructions Governing Train Operation at Station；
        technologic instructions of station
        简称《站细》，是根据《技规》和《行规》的有关规定，结合车站技术设备和作业等具体条件所制定的车站技术管理和作业组织的基本技术文件。
       """;
    Response<Embedding> embeddingResponse = embeddingModel.embed(prompt);
    log.info(embeddingResponse.toString());
    return embeddingResponse.content().toString();
  }

  /**
   * 新建向量数据库实例和创建索引：test-qdrant 类似mysql create database test-qdrant
   * http://localhost:9012/embedding/createCollection
   */
  @GetMapping(value = "/embedding/createCollection")
  public void createCollection() {
    var vectorParams =
        Collections.VectorParams.newBuilder()
            .setDistance(Collections.Distance.Cosine)
            .setSize(1024)
            .build();
    qdrantClient.createCollectionAsync(QdrantConstant.COLLECTION_NAME, vectorParams);
  }

  /**
   * 往向量数据库新增文本记录
   *
   * @return the string
   */
  @GetMapping(value = "/embedding/add")
  public String add() {
    String prompt =
        """
            1.铁路行车组织
            organization of train operation
            铁路行车组织是铁路运输组织的重要组成部分，是铁路综合运用各种技术设备合理组织列车运行以实现旅客和货物运输过程的生产计划与组织工作。它的主内容包括：车站技术作业组织、车流组织、列车运行图和铁路通过能力、铁路运输工作技术计划和调度工作等。
            2.铁路技术管理规程
            Regulations Governing Railway Technical Operation
            简称《技规》，是铁路进行技术管理和从事运输生产的基本法规。它的主要内容包括：铁路技术设备的基本要求和标准，行车组织工作应遵循的基本原则、工作方法和作业程序、信号显示的要求和方法以及铁路运输工作人员的主要职责和必须具备的基本条件。
            3.铁路行车组织规则
            Rules Governing Organization of Train Operation；operating rules
            简称《行规》，是铁路局根据《技规》中行车组织篇规定的原则，结合铁路局管内技术设备的具体条件所制定的一种行车组织的补充规则，作为《技规》的补充。
            4.车站行车工作细则
            Detailed Instructions Governing Train Operation at Station；
            technologic instructions of station
            简称《站细》，是根据《技规》和《行规》的有关规定，结合车站技术设备和作业等具体条件所制定的车站技术管理和作业组织的基本技术文件。
           """;
    TextSegment segment1 = TextSegment.from(prompt);
    segment1.metadata().put("name", "铁路术语");
    Embedding embedding1 = embeddingModel.embed(segment1).content();
    String result = embeddingStore.add(embedding1, segment1);

    System.out.println(result);

    return result;
  }

  @GetMapping(value = "/embedding/query1")
  public String query1() {
    Embedding queryEmbedding = embeddingModel.embed("什么是站细").content();
    EmbeddingSearchRequest embeddingSearchRequest =
        EmbeddingSearchRequest.builder().queryEmbedding(queryEmbedding).maxResults(1).build();
    EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);
    String text = searchResult.matches().getFirst().embedded().text();
    log.info(text);
    return text;
  }

  @GetMapping(value = "/embedding/query2")
  public String query2() {
    Embedding queryEmbedding = embeddingModel.embed("站细").content();

    EmbeddingSearchRequest embeddingSearchRequest =
        EmbeddingSearchRequest.builder()
            .queryEmbedding(queryEmbedding)
            .filter(metadataKey("name").isEqualTo("铁路术语"))
            .maxResults(1)
            .build();

    EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);
    String text = searchResult.matches().getFirst().embedded().text();
    log.info(text);
    return text;
  }
}
